{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9fd54a32",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/upstage.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3a8796-edc8-43f2-94ad-fe4fb20d70ed",
   "metadata": {},
   "source": [
    "# Upstage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "081d07d2",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f6f8702",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-upstage llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b007403c-6b7a-420c-92f1-4171d05ed9bb",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ead155e-b8bd-46f9-ab9b-28fc009361dd",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "859f7447-bbc3-4036-b541-a90854a732f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"UPSTAGE_API_KEY\"] = \"YOUR_API_KEY\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60be18ae-c957-4ac2-a58a-0652e18ee6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.upstage import Upstage\n",
    "\n",
    "llm = Upstage(\n",
    "    model=\"solar-mini\",\n",
    "    # api_key=\"YOUR_API_KEY\"  # uses UPSTAGE_API_KEY env var by default\n",
    ")\n",
    "\n",
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac2cbebb-a444-4a46-9d85-b265a3483d68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the venture capital firm Y Combinator, which has funded and incubated many successful startups. He is also the author of several influential essays on entrepreneurship, startup culture, and technology.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14831268-f90f-499d-9d86-925dbc88292b",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe29574-4af1-48d5-9739-f60652b6ce6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "resp = llm.chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cbd550a-0264-4a11-9b2c-a08d8723a5ae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: I am Captain Redbeard, the fearless pirate!\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ed5e894-4597-4911-a623-591560f72b82",
   "metadata": {},
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb7986f-aaed-42e2-abdd-f274f6d4fc59",
   "metadata": {},
   "source": [
    "Using `stream_complete` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d43f17a2-0aeb-464b-a7a7-732ba5e8ef24",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.stream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0214e911-cf0d-489c-bc48-9bb1d8bf65d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known for co-founding the startup accelerator Y Combinator, which has helped launch some of the most successful tech companies in the world, including Airbnb, Dropbox, and Stripe. He is also the author of several influential essays on startup culture and entrepreneurship, including \"How to Start a Startup\" and \"Hackers & Painters.\""
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40350dd8-3f50-4a2f-8545-5723942039bb",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc636e65-a67b-4dcd-ac60-b25abc9d8dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4475a6bc-1051-4287-abce-ba83324aeb9e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I am Captain Redbeard, the fearless pirate!"
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90f07f7e-927f-47a2-9797-de5a86d61e1f",
   "metadata": {},
   "source": [
    "## Function Calling\n",
    "\n",
    "Upstage models have native support for function calling. This conveniently integrates with LlamaIndex tool abstractions, letting you plug in any arbitrary Python function to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "212bb2d2-2bed-4188-85ad-3cd497d4b864",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "class Song(BaseModel):\n",
    "    \"\"\"A song with name and artist\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "\n",
    "\n",
    "def generate_song(name: str, artist: str) -> Song:\n",
    "    \"\"\"Generates a song with provided name and artist.\"\"\"\n",
    "    return Song(name=name, artist=artist)\n",
    "\n",
    "\n",
    "tool = FunctionTool.from_defaults(fn=generate_song)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba29ea39-39cd-42f4-934e-78fba4935a1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='My Song' artist='John Doe'\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.upstage import Upstage\n",
    "\n",
    "llm = Upstage()\n",
    "response = llm.predict_and_call([tool], \"Generate a song\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "629bdb64-e5bd-4539-9537-54352dd8dbb4",
   "metadata": {},
   "source": [
    "We can also do multiple function calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c098cb38-bfed-4907-8109-6810756f2ae5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: generate_song, Input: {'args': (), 'kwargs': {'name': 'Beatles', 'artist': 'Beatles'}}, Output: name='Beatles' artist='Beatles'\n"
     ]
    }
   ],
   "source": [
    "llm = Upstage()\n",
    "response = llm.predict_and_call(\n",
    "    [tool],\n",
    "    \"Generate five songs from the Beatles\",\n",
    "    allow_parallel_tool_calls=True,\n",
    ")\n",
    "for s in response.sources:\n",
    "    print(f\"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df5fa1ab-f598-46da-80f3-f6af5dd2fe83",
   "metadata": {},
   "source": [
    "## Async"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e088b90-12b6-4211-a9ca-696e9897e9ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.upstage import Upstage\n",
    "\n",
    "llm = Upstage()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f649683-896a-439e-b14b-e5df5d803b82",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = await llm.acomplete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d05d8b7c-07b3-4ce8-9a8c-fcd1e830316c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the startup accelerator Y Combinator, which has helped launch and fund many successful tech companies. He is also the author of several influential essays on startups, entrepreneurship, and technology, including \"How to Start a Startup\" and \"Hackers & Painters.\"\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc72d09-3bcd-4d48-9bb7-0920c56310c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = await llm.astream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b28d7a06-1518-4ec0-b3cc-6364395b3561",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a computer scientist, entrepreneur, and essayist. He is best known as the co-founder of the startup accelerator Y Combinator, which has helped launch some of the most successful tech companies in the world, including Airbnb, Dropbox, and Stripe. Graham is also a prolific writer, and his essays on topics such as startup advice, artificial intelligence, and the future of work have been widely read and influential in the tech industry."
     ]
    }
   ],
   "source": [
    "async for delta in resp:\n",
    "    print(delta.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3782e752-3c65-4b11-b2c6-8efe55fc33b1",
   "metadata": {},
   "source": [
    "Async function calling is also supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dcfaa77-edfe-42f9-8138-884c99dd3e62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='My Song' artist='Me'\n"
     ]
    }
   ],
   "source": [
    "llm = Upstage()\n",
    "response = await llm.apredict_and_call([tool], \"Generate a song\")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
